import numpy as np

# 1. 生成数据
def generate_vectors(num_vectors=100, dimensions=3):
    """ 
    生成向量数据
    :param num_vectors: 向量数量
    :param dimensions: 向量维度
    :return: 随机生成的向量数据
    """ 
    np.random.seed(42)
    return np.random.rand(num_vectors, dimensions)

data_vectors = generate_vectors(num_vectors=100, dimensions=3)

# 2. LSH实现
class LSH:
    """ 
    基于余弦相似度的LSH实现，支持参数调优
    """ 
    def __init__(self, dimensions, num_hashes, num_tables):
        self.dimensions = dimensions
        self.num_hashes = num_hashes
        self.num_tables = num_tables
        self.hash_planes = [np.random.randn(num_hashes, dimensions) for _ in range(num_tables)]  # 每个表随机生成超平面
        self.tables = [{} for _ in range(num_tables)]  # 修正：使用列表推导式创建字典列表

    def hash_function(self, vector, planes):
        """ 
        对单个向量进行哈希
        :param vector: 输入向量
        :param planes: 超平面集合
        :return: 哈希值
        """ 
        projections = np.dot(planes, vector)
        return ''.join(['1' if p > 0 else '0' for p in projections])

    def insert(self, vectors):
        """
        将向量插入所有哈希表中
        :param vectors: 输入向量列表
        """
        for table_id, planes in enumerate(self.hash_planes):
            for idx, vector in enumerate(vectors):
                hash_value = self.hash_function(vector, planes)
                if hash_value not in self.tables[table_id]:
                    self.tables[table_id][hash_value] = [] 
                self.tables[table_id][hash_value].append(idx)

    def query(self, query_vector):
        """
        查询向量
        :param query_vector: 查询向量
        :return: 所有匹配的向量索引
        """
        candidates = set()
        for table_id, planes in enumerate(self.hash_planes):
            hash_value = self.hash_function(query_vector, planes)
            if hash_value in self.tables[table_id]:
                candidates.update(self.tables[table_id][hash_value])
        return list(candidates)

# 3. 初始化LSH并进行相似性搜索
num_hashes = 5
num_tables = 3

lsh = LSH(dimensions=3, num_hashes=num_hashes, num_tables=num_tables)
lsh.insert(data_vectors)

# 4. 查询向量
query_vector = np.array([0.5, 0.5, 0.5])  # 查询向量
result_indices = lsh.query(query_vector)

# 5. 参数调优分析
def analyze_lsh_performance(vectors, lsh, query_vector):
    """  
    分析LSH性能  
    :param vectors: 输入向量  
    :param lsh: LSH 实例  
    :param query_vector: 查询向量  
    """  
    result_indices = lsh.query(query_vector)  
    print(f"查询向量: {query_vector}")  
    print(f"匹配向量索引: {result_indices}")  
    print(f"匹配向量: {[vectors[idx] for idx in result_indices]}")

analyze_lsh_performance(data_vectors, lsh, query_vector)

# 输出哈希表信息
print("\n哈希表分布:")


for table_id, table in enumerate(lsh.tables):
    print(f"哈希表 {table_id}: {len(table)} 桶")